{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练量子隐形传态协议\n",
    "\n",
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "量子隐形传态（quantum teleportation）是可以通过本地操作和经典通信（LOCC）协议完成的另一项重要任务，该协议借助提前制备好的纠缠资源在两个空间上分离的通信节点（仅允许经典信道）之间传输量子信息。在本教程中，我们将首先简要回顾一下量子隐形传态协议，并使用量桨进行模拟。然后，我们将介绍如何使用 LOCCNet 学习出一个量子隐形传态协议。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 量子隐形传态协议\n",
    "\n",
    "量子隐形传态协议最初由 C. H. Bennett 等人在 1993 年提出 [1]，并在 1997 年通过基于光子的实验平台进行了验证 [2-3]。其主要工作流程如下图所示。传输过程中需要 2 个通信节点或参与方，即 Alice 和 Bob。简单起见，我们仅考虑传输一个单量子比特的量子态 $|\\psi\\rangle_C$，这样整个系统总共需要 3 个量子比特，包括提前共享给 Alice 和 Bob 的最大纠缠态 $|\\Phi^+\\rangle_{AB} = \\frac{1}{\\sqrt{2}}(|00\\rangle + |11\\rangle)$。 Alice 持有量子比特 A 和 C， Bob 持有量子比特 B。**注释：量子隐形传态协议仅传输量子信息，而非在物理上直接传输量子比特。**\n",
    "\n",
    "![](figures/teleportation-fig-circuit.jpg \"图 1：量子隐形传态: 传输量子态 $|\\psi\\rangle$ 从 Alice 到 Bob\")\n",
    "<div style=\"text-align:center\">图 1：量子隐形传态: 传输量子态 $|\\psi\\rangle$ 从 Alice 到 Bob </div>\n",
    "\n",
    "**步骤 I：** 一开始，整个量子系统的状态可以描述为\n",
    "\n",
    "$$\n",
    "\\lvert\\varphi_{0}\\rangle \n",
    "= \\lvert\\psi\\rangle_{C}\\otimes \\lvert\\Phi^+\\rangle_{AB} \n",
    "= \\frac{1}{\\sqrt{2}}\\big[\\alpha\\lvert0\\rangle(\\lvert00\\rangle + \\lvert11\\rangle)+\\beta\\lvert1\\rangle(\\lvert00\\rangle + \\lvert11\\rangle)\\big],\n",
    "\\tag{1}\n",
    "$$\n",
    "\n",
    " 其中 Alice 要传输的量子态是 $|\\psi\\rangle_C = \\alpha|0\\rangle_C + \\beta|1\\rangle_C$ 并且 $\\alpha, \\beta \\in \\mathbb{C}$。\n",
    "\n",
    "**步骤 II：**  Alice 在她持有的两个量子比特上作用 CNOT 门，得到作用后的量子态为\n",
    "\n",
    "$$\n",
    "|\\varphi_1\\rangle  \n",
    "= \\frac{1}{\\sqrt{2}}\\big[\\alpha\\lvert0\\rangle(\\lvert00\\rangle + \\lvert11\\rangle)+\\beta\\lvert1\\rangle(\\lvert10\\rangle + \\lvert01\\rangle)\\big],\n",
    "\\tag{2}\n",
    "$$\n",
    "\n",
    "**步骤 III：** 接着 Alice 在她持有的量子比特 $C$ 上作用 Hadamard 门，使得整个系统的状态变为 $|\\varphi_2\\rangle$\n",
    "\n",
    "$$\n",
    "|\\varphi_2\\rangle = \\frac{1}{2}\\big[\\alpha(\\lvert0\\rangle + \\lvert1\\rangle)(\\lvert00\\rangle + \\lvert11\\rangle)+\\beta(\\lvert0\\rangle - \\lvert1\\rangle)(\\lvert10\\rangle + \\lvert01\\rangle)\\big],\n",
    "\\tag{3}\n",
    "$$\n",
    "\n",
    "为了更好地观察后续的测量结果，不妨将上述态重新写为\n",
    "\n",
    "$$\n",
    "\\lvert\\varphi_{2}\\rangle = \\frac{1}{2}\\big[\\lvert00\\rangle(\\alpha\\lvert0\\rangle + \\beta\\lvert1\\rangle) + \\lvert01\\rangle(\\alpha\\lvert1\\rangle + \\beta\\lvert0\\rangle) + \\lvert10\\rangle(\\alpha\\lvert0\\rangle - \\beta\\lvert1\\rangle) + \\lvert11\\rangle(\\alpha\\lvert1\\rangle - \\beta\\lvert0\\rangle)\\big].\n",
    "\\tag{4}\n",
    "$$\n",
    "\n",
    "**步骤 IV：**  Alice 在计算基（computational basis）$\\{|00\\rangle, |01\\rangle, |10\\rangle, |11\\rangle\\}$ 上测量她的两个量子比特并将结果 $m_1m_2$ 通过经典信道发送给 Bob 。 总共有 4 种不同的可能性： $m_1m_2 \\in \\{ 00, 01,10, 11\\}$。 然后， Bob 根据收到的消息在其量子比特 $B$ 上进行对应的操作。\n",
    "\n",
    "- 如果测量结果为 $m_1m_2 = 00$，则 Bob 的态将为 $\\alpha\\lvert0\\rangle + \\beta\\lvert1\\rangle$。 无需任何操作即可完成传输。\n",
    "- 如果测量结果为 $ m_1m_2 = 01 $，则 Bob 的态将为 $\\alpha\\lvert1\\rangle + \\beta\\lvert0\\rangle$。  Bob 需要在其量子比特上作用 $ X $ 门。\n",
    "- 如果测量结果为 $ m_1m_2 = 10 $，则 Bob 的态将为 $\\alpha\\lvert0\\rangle - \\beta\\lvert1\\rangle$。  Bob 需要在其量子比特上作用 $ Z $ 门。\n",
    "- 如果测量结果为 $ m_1m_2 = 11 $，则 Bob 的态将为 $\\alpha\\lvert1\\rangle - \\beta\\lvert0\\rangle$。 Bob 需要在其量子比特上执行 $ X $ 门操作，然后执行 $ Z $ 门操作。\n",
    "\n",
    "在下一节中，我们将介绍如何使用量桨模拟量子隐形传态协议。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paddle Quantum 代码实现\n",
    "\n",
    "首先，我们需要导入所有依赖包："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:22:58.705389Z",
     "start_time": "2021-03-09T06:22:55.030541Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from paddle_quantum.locc import LoccNet\n",
    "from paddle import matmul, trace\n",
    "import paddle\n",
    "from paddle_quantum.utils import state_fidelity\n",
    "from paddle_quantum.state import bell_state, isotropic_state, density_op_random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "初始化整个量子系统，然后定义量子电路和隐形传态协议。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:23:06.232030Z",
     "start_time": "2021-03-09T06:23:06.179323Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC, self).__init__()\n",
    "        \n",
    "        # 添加第一个参与方 Alice\n",
    "        # 第一个参数 2 代表着 Alice 手里有几个量子比特\n",
    "        # 第二个参数代表着参与方的名字\n",
    "        self.add_new_party(2, party_name=\"Alice\")\n",
    "        \n",
    "        # 添加第二个参与方 Bob\n",
    "        # 第一个参数 1 代表着 Bob 手里有几个量子比特\n",
    "        # 第二个参数代表着参与方的名字\n",
    "        self.add_new_party(1, party_name=\"Bob\")\n",
    "        \n",
    "\n",
    "        # 准备一个贝尔态\n",
    "        _state = paddle.to_tensor(bell_state(2))\n",
    "        # _state = paddle.to_tensor(isotropic_state(2, 0.8))\n",
    "\n",
    "        \n",
    "        # 随机制备传输用的纯态 （rank =1）\n",
    "        random_state = density_op_random(n=1, real_or_complex=2, rank=1)\n",
    "        self.state_C = paddle.to_tensor(random_state)\n",
    "        \n",
    "        # 通过分配上述制备好的量子态初始化整个量子系统\n",
    "        # 这里 (\"Alice\", 0) 即表示量子比特 C\n",
    "        # 这里 (\"Alice\", 1) 即表示量子比特 A\n",
    "        # 这里 (\"Bob\", 0) 即表示量子比特 B\n",
    "        # print('提前分配好的纠缠态为:\\n', _state.numpy())\n",
    "        self.set_init_state(self.state_C, [(\"Alice\", 0)])\n",
    "        self.set_init_state(_state, [(\"Alice\", 1), (\"Bob\", 0)])\n",
    "\n",
    "\n",
    "    def teleportation(self):\n",
    "        status = self.init_status\n",
    "        \n",
    "        # 设置 Alice 的本地操作\n",
    "        cirA = self.create_ansatz(\"Alice\")\n",
    "        cirA.cnot([0, 1])\n",
    "        cirA.h(0)\n",
    "        \n",
    "        # 运行上述电路\n",
    "        status = cirA.run(status)\n",
    "        \n",
    "        # Alice 在计算基上测量她所持有的两个量子比特 C 还有 A\n",
    "        # 得到并记录四种结果 00，01，10，11\n",
    "        status_A = self.measure(status, [(\"Alice\", 0), (\"Alice\", 1)], [\"00\", \"01\", \"10\", \"11\"])\n",
    "        \n",
    "        # 用于记录平均保真度\n",
    "        fid_list = []\n",
    "\n",
    "        # Bob 根据 Alice 的测量结果选择不同的门作用在自己的量子比特上\n",
    "        for i, s in enumerate(status_A):\n",
    "            \n",
    "            # 判断语句根据 Alice 的测量结果，进行不同操作 \n",
    "            if status_A[i].measured_result == '00':\n",
    "                \n",
    "                # 创建 Bob 的本地操作\n",
    "                cirB = self.create_ansatz(\"Bob\")\n",
    "                \n",
    "                # 执行电路\n",
    "                status_B = cirB.run(s) \n",
    "                \n",
    "                # 仅保留 Bob 的量子比特 B\n",
    "                status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "\n",
    "                # 计算初始态和传输后态之间的保真度\n",
    "                fid = state_fidelity(self.state_C.numpy(), status_fin.state.numpy())**2\n",
    "                fid_list.append(fid * status_fin.prob.numpy()[0])\n",
    "                \n",
    "            # 以下操作类似\n",
    "            elif status_A[i].measured_result == '01':\n",
    "                cirB = self.create_ansatz(\"Bob\")\n",
    "                cirB.x(0)\n",
    "                status_B = cirB.run(s)\n",
    "                status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "                fid = state_fidelity(self.state_C.numpy(), status_fin.state.numpy())**2\n",
    "                fid_list.append(fid * status_fin.prob.numpy()[0])\n",
    "            elif status_A[i].measured_result == '10':\n",
    "                cirB = self.create_ansatz(\"Bob\")\n",
    "                cirB.z(0)\n",
    "                status_B = cirB.run(s)\n",
    "                status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "                fid = state_fidelity(self.state_C.numpy(), status_fin.state.numpy())**2\n",
    "                fid_list.append(fid * status_fin.prob.numpy()[0])\n",
    "            elif status_A[i].measured_result == '11':\n",
    "                cirB = self.create_ansatz(\"Bob\")\n",
    "                cirB.x(0)\n",
    "                cirB.z(0)\n",
    "                status_B = cirB.run(s)\n",
    "                status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "                fid = state_fidelity(self.state_C.numpy(), status_fin.state.numpy())**2\n",
    "                fid_list.append(fid * status_fin.prob.numpy()[0])\n",
    "        fid_avg = sum(fid_list)\n",
    "        \n",
    "        return fid_avg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后，我们随机生成 200 个量子纯态，并使用态保真度 $ F $ 来衡量传输协议好坏，其中\n",
    "\n",
    "$$\n",
    "F(\\rho,\\sigma) \\equiv \\text{tr}\\big( \\sqrt{\\sqrt{\\rho}\\sigma \\sqrt{\\rho}} \\big)^2.\n",
    "\\tag{5}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:23:14.626278Z",
     "start_time": "2021-03-09T06:23:08.646876Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "平均保真度 = 1.0 , 标准差 = 6.429150606796583e-09\n"
     ]
    }
   ],
   "source": [
    "SEED = 999              # 固定随机数\n",
    "num_state = 200         # 设置随机态的生成数量\n",
    "list_fid = []           # 用于记录保真度\n",
    "np.random.seed(SEED)\n",
    "# 开始采样 \n",
    "for idx in range(num_state):\n",
    "    list_fid.append(LOCC().teleportation())\n",
    "\n",
    "print('平均保真度 =', np.around(sum(list_fid)/len(list_fid), 4), ', 标准差 =', np.std(list_fid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**注释：** 我们要指出的是，该协议的有效性取决于提前分配好的贝尔态的品质。 感兴趣的读者可以将纠缠态从 `bell_state(2)` 更改为 `isotropic_state(2, p)`，并测试观察在贝尔态中出现的量子噪声将如何影响该协议的性能。 其中 isotropic 态的定义如下，\n",
    "\n",
    "$$\n",
    "\\rho_{\\text{iso}}(p) = p\\lvert\\Phi^+\\rangle \\langle\\Phi^+\\rvert + (1-p)\\frac{I}{4}, \\quad p \\in [0,1]\n",
    "\\tag{6}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 LOCCNet 学习量子隐形传态\n",
    "\n",
    "### 训练一个自定义的 LOCC 协议\n",
    "\n",
    "一般的 LOCC 协议可以通过经典通信的回合数 $ r $ 来进行分类。 原始的量子隐形传态协议是单轮通讯协议（$ r = 1 $）。 为简单起见，我们也将通信回合数限制为单轮。与原始的协议不同，我们将使用参数化量子电路（PQC）来把 Bob 作用在量子比特上的固定门 $U\\in\\{X,Z\\}$ 替换成布洛赫球面上的一个广义旋转门 $ U_3 $， 其定义为\n",
    "\n",
    "$$\n",
    "U_3(\\theta, \\phi, \\varphi) =\n",
    "\\begin{bmatrix}\n",
    "\\cos(\\frac{\\theta}{2})           & -e^{i\\varphi}\\sin(\\frac{\\theta}{2})\\\\\n",
    "e^{i\\phi}\\sin(\\frac{\\theta}{2})  & e^{i(\\phi+\\varphi)} \\cos(\\frac{\\theta}{2})\n",
    "\\end{bmatrix}.\n",
    "\\tag{7}\n",
    "$$\n",
    "\n",
    "这将为我们带来更强大的搜索能力，来寻找更多使用场景下的 LOCC 协议。类似地，我们将 Alice 的本地操作更改为更通用的参数化量子电路，这里具体使用的是一个内置的两量子比特通用门 `universal_2_qubit_gate(theta, which_qubit=[0,1])` [4]。基于以上条件，我们训练一个自定义的 LOCC 协议流程如下：\n",
    "\n",
    "1. Alice 对她所持有的两个量子比特作用两量子比特通用门。\n",
    "2. 然后 Alice 在计算基上测量她的两个量子比特，并通过经典信道与 Bob 交流。\n",
    "3. 共计有 4 种可能的测量结果：$m_1m_2 \\in \\{00, 01, 10, 11\\}$。 Bob 需要根据这些测量结果采取不同的本地操作。在 Bob 进行操作后，记其量子态为 $\\lvert \\psi\\rangle_{B}$。\n",
    "4. 计算 $\\lvert \\psi\\rangle_{B}$ 与 $\\lvert\\psi\\rangle_C$（纯态）之间的量子态重叠（state overlap）并记为 $ O $。 由于 LOCCNet 框架目前仅支持密度矩阵形式，因此我们必须将它们重写为 $\\rho_{B} = |\\psi\\rangle\\langle\\psi|_B$ 和 $\\rho_{C} = |\\psi\\rangle\\langle\\psi|_C$。然后可以得到 $O = \\text{Tr}(\\rho_C\\rho_{B})$。对于纯态，此距离度量就是保真度 (fidelity)。\n",
    "5. 将损失函数设置为 4 种可能测量结果的累加，即 $L = \\sum_{m_1m_2} \\big(1-\\text{Tr}(\\rho_C\\rho_{B})\\big)$，并使用基于梯度的优化方法更新 Alice 和 Bob 本地操作中的参数，从而使得损失函数最小化。\n",
    "6. 重复步骤 1-5，直到损失函数收敛。\n",
    "7. 生成一组随机的态 $\\{\\lvert\\psi_C\\rangle\\}$，并以平均保真度对训练出的传输协议进行基准测试。\n",
    "\n",
    "![](figures/teleportation-fig-LOCCNet.png \"图 2： 用 LOCCNet 学习量子隐形传态协议\")\n",
    "<div style=\"text-align:center\">图 2： 用 LOCCNet 学习量子隐形传态协议 </div>\n",
    "\n",
    "**注释：** 为了确保训练出的 LOCC 协议对所有态均有效，我们将训练集设为 4 个线性独立态，即 $\\{|0\\rangle\\langle 0|,|1\\rangle\\langle 1|,|+\\rangle\\langle +|,|+\\rangle\\langle +|_y\\}$ 。其中 $|+\\rangle, |+\\rangle_y$ 分别为泡利 $X,Y$ 矩阵的正特征值对应的特征向量。 该训练集在密度矩阵的表示下为：\n",
    "\n",
    "$$\n",
    "\\rho_0 = \\left[\\begin{array}{ccc}\n",
    "1 & 0\\\\\n",
    "0 & 0\n",
    "\\end{array}\\right], \n",
    "\\rho_1 = \\left[\\begin{array}{ccc}\n",
    "0 & 0\\\\\n",
    "0 & 1\n",
    "\\end{array}\\right], \n",
    "\\rho_2 = \\left[\\begin{array}{ccc}\n",
    "0.5 & 0.5\\\\\n",
    "0.5 & 0.5\n",
    "\\end{array}\\right], \n",
    "\\rho_3 = \\left[\\begin{array}{ccc}\n",
    "0.5 & -0.5 i\\\\\n",
    "0.5i & 0.5\n",
    "\\end{array}\\right]. \n",
    "\\tag{8}\n",
    "$$\n",
    "\n",
    "任何一个单量子比特量子态都可以写为上述 4 个态的线性组合。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:26:15.991221Z",
     "start_time": "2021-03-09T06:26:15.937457Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC_Train(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC_Train, self).__init__()\n",
    "        \n",
    "        # 初始化 LOCCNet\n",
    "        self.parties = list()\n",
    "        \n",
    "        # 添加第一个参与方 Alice\n",
    "        # 第一个参数 2 代表着 Alice 手里有几个量子比特\n",
    "        # 第二个参数代表着参与方的名字\n",
    "        self.add_new_party(2, party_name=\"Alice\")\n",
    "        \n",
    "        # 添加第二个参与方 Bob\n",
    "        # 第一个参数 1 代表着 Bob 手里有几个量子比特\n",
    "        # 第二个参数代表着参与方的名字\n",
    "        self.add_new_party(1, party_name=\"Bob\")\n",
    "        \n",
    "\n",
    "        # 准备一个贝尔态\n",
    "        _state = paddle.to_tensor(bell_state(2))\n",
    "        # _state = paddle.to_tensor(isotropic_state(2, 0.8))\n",
    "\n",
    "\n",
    "        # 设置 Alice 本地操作的参数\n",
    "        self.theta_A = self.create_parameter(shape=[15], default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2 * np.pi), dtype=\"float64\")\n",
    "        # 设置 Bob 本地操作的参数\n",
    "        self.theta_B = self.create_parameter(shape=[4, 3], default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2 * np.pi), dtype=\"float64\")\n",
    "\n",
    "        # 训练集: 4 个线性独立态\n",
    "        _state0 = paddle.to_tensor(np.array([[1, 0], [0, 0]], dtype=np.complex128))\n",
    "        _state1 = paddle.to_tensor(np.array([[0, 0], [0, 1]], dtype=np.complex128))\n",
    "        _state2 = paddle.to_tensor(np.array([[0.5, 0.5], [0.5, 0.5]], dtype=np.complex128))\n",
    "        _state3 = paddle.to_tensor(np.array([[0.5, -0.5j], [0.5j, 0.5]], dtype=np.complex128))\n",
    "        self.init_states = [_state0, _state1, _state2, _state3]\n",
    "\n",
    "        # 通过分配上述制备好的量子态初始化整个量子系统\n",
    "        self.set_init_state(_state, [(\"Alice\", 1), (\"Bob\", 0)])\n",
    "        self.set_init_state(_state0, [(\"Alice\", 0)])\n",
    "\n",
    "    def LOCCNet(self):\n",
    "        \n",
    "        # 定义训练过程\n",
    "        loss = 0\n",
    "        temp_state = self.init_status\n",
    "        \n",
    "        # 开始训练\n",
    "        for init_state in self.init_states:\n",
    "            \n",
    "            # 重置 Alice 持有的量子比特 C 至训练集中的量子态\n",
    "            status = self.reset_state(temp_state, init_state, [(\"Alice\", 0)])\n",
    "\n",
    "            # 定义 Alice 的本地操作\n",
    "            cirA = self.create_ansatz(\"Alice\")\n",
    "            cirA.universal_2_qubit_gate(self.theta_A, [0, 1])\n",
    "\n",
    "            # 执行 Alice 的电路\n",
    "            status = cirA.run(status)\n",
    "            \n",
    "             # 测量得到四个可能的结果\n",
    "            status_A = self.measure(status, [(\"Alice\", 0), (\"Alice\", 1)], [\"00\", \"01\", \"10\", \"11\"])\n",
    "            \n",
    "            # Bob 根据测量结果选择不同的门作用在自己的量子比特上\n",
    "            for i, s in enumerate(status_A):\n",
    "                \n",
    "                # 定义 Bob 的本地操作\n",
    "                cirB = self.create_ansatz(\"Bob\")\n",
    "                \n",
    "                # 作用单量子比特通用门\n",
    "                cirB.u3(*self.theta_B[i], 0)\n",
    "                \n",
    "                # 执行 Bob 的电路\n",
    "                status_B = cirB.run(s)\n",
    "                \n",
    "                # 仅留下 Bob 的量子比特 B\n",
    "                status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "                \n",
    "                # 将所有的测量结果的损失函数进行累加\n",
    "                loss += 1 - paddle.real(trace(matmul(init_state, status_fin.state)))[0]\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    # 存储训练结束后的最优参数\n",
    "    def save_module(self):\n",
    "        theta = np.array([self.theta_A.numpy(), self.theta_B.numpy()], dtype=\"object\")\n",
    "        np.save('parameters/QT_LOCCNet', theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:27:21.094741Z",
     "start_time": "2021-03-09T06:26:16.291397Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "itr 0: 7.721435597307273\n",
      "itr 10: 0.5636061774837096\n",
      "itr 20: 0.20642802374795632\n",
      "itr 30: 0.07340371172539029\n",
      "itr 40: 0.025112969166681598\n",
      "itr 50: 0.008655519545057166\n",
      "itr 60: 0.0035853096929094885\n",
      "itr 70: 0.0014441032202451298\n",
      "itr 80: 0.0004540960989092291\n",
      "itr 90: 0.00018530025729923683\n",
      "itr 100: 4.679932743989479e-05\n",
      "itr 110: 1.5186530775368468e-05\n",
      "itr 120: 4.8322085606233856e-06\n",
      "itr 130: 3.08099659207528e-06\n",
      "itr 140: 1.0481638913484304e-06\n"
     ]
    }
   ],
   "source": [
    "ITR = 150   # 设置优化循环次数\n",
    "LR = 0.2    # 设置学习速率\n",
    "SEED = 999  # 固定本地操作中参数的初始化随机数种子\n",
    "np.random.seed(SEED)\n",
    "paddle.seed(SEED)\n",
    "\n",
    "\n",
    "net = LOCC_Train()\n",
    "\n",
    "# 选择 Adam 优化器\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters= net.parameters())\n",
    "\n",
    "# 记录损失\n",
    "loss_list = []\n",
    "\n",
    "# 开始优化循环\n",
    "for itr in range(ITR):\n",
    "\n",
    "    # 向前传播计算损失函数\n",
    "    loss = net.LOCCNet()\n",
    "\n",
    "    # 反向传播优化损失函数\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "\n",
    "    # 清除梯度\n",
    "    opt.clear_grad()\n",
    "\n",
    "    # 记录训练结果\n",
    "    loss_list.append(loss.numpy()[0])\n",
    "    if itr % 10 == 0:\n",
    "        print(\"itr \" + str(itr) + \":\", loss.numpy()[0])\n",
    "\n",
    "# 保存参数\n",
    "net.save_module()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基准测试\n",
    "\n",
    "如果想要更快捷的感受效果，我们还提供**加载预先训练好的电路参数**并直接测试性能。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:27:35.419952Z",
     "start_time": "2021-03-09T06:27:35.391435Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC_Test(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC_Test, self).__init__()\n",
    "        \n",
    "        self.parties = list()\n",
    "        self.add_new_party(2, party_name=\"Alice\")\n",
    "        self.add_new_party(1, party_name=\"Bob\")\n",
    "        \n",
    "        self.theta_A = paddle.to_tensor(para[0])\n",
    "        self.theta_B = paddle.to_tensor(para[1])\n",
    "        \n",
    "        _state = paddle.to_tensor(bell_state(2))\n",
    "        random_state = density_op_random(n=1)\n",
    "        self._state0 = paddle.to_tensor(random_state)\n",
    "        self.set_init_state(_state, [(\"Alice\", 1), (\"Bob\", 0)])\n",
    "        self.set_init_state(self._state0, [(\"Alice\", 0)])\n",
    "        \n",
    "\n",
    "    def benchmark(self):\n",
    "        input_state = self.init_status\n",
    "        cirA = self.create_ansatz(\"Alice\")\n",
    "        cirA.universal_2_qubit_gate(self.theta_A, [0, 1])\n",
    "\n",
    "        status = cirA.run(input_state)\n",
    "        status_A = self.measure(status, [(\"Alice\", 0), (\"Alice\", 1)], [\"00\", \"01\", \"10\", \"11\"])\n",
    "        fid_list = []\n",
    "\n",
    "        for i, s in enumerate(status_A):\n",
    "            cirB = self.create_ansatz(\"Bob\")\n",
    "            cirB.u3(*self.theta_B[i], 0)\n",
    "            status_B = cirB.run(s)\n",
    "            status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "            fid = state_fidelity(self._state0.numpy(), status_fin.state.numpy())**2\n",
    "            fid_list.append(fid*status_fin.prob.numpy()[0])\n",
    "        fid_ave = sum(fid_list)\n",
    "        \n",
    "        return fid_ave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:28:31.015666Z",
     "start_time": "2021-03-09T06:28:18.203521Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "平均保真度 = 1.0 , 标准差 = 5.695904176663085e-07\n"
     ]
    }
   ],
   "source": [
    "# 加载预先训练的电路参数\n",
    "para = np.load('parameters/QT_LOCCNet_Trained.npy', allow_pickle=True)\n",
    "SEED = 999              # 固定生成传输态的随机种子\n",
    "num_state = 200         # 设置随机态的生成数量\n",
    "list_fid = []           # 用于记录保真度\n",
    "np.random.seed(SEED)\n",
    "paddle.seed(SEED)\n",
    "\n",
    "# 采样\n",
    "for idx in range(num_state):\n",
    "    list_fid.append(LOCC_Test().benchmark())\n",
    "\n",
    "print('平均保真度 =', np.around(sum(list_fid)/len(list_fid), 4), ', 标准差 =', np.std(list_fid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 结论\n",
    "\n",
    "从以上数值实验的结果可以看出，基于 LOCCNet 我们成功地学习出了量子隐形传态协议。 最初的隐形传态协议旨在传输单量子比特量子态，无法直接推广到多量子比特的情形。 相比之下，LOCCNet 为寻找多量子比特情形下的隐形传态协议提供了可能, 更是可以训练得到更一般的通过量子纠缠和 LOCC 实现不同节点间量子信道的模拟。 此外，感兴趣的读者可以尝试各类量子噪声会如何影响 LOCCNet 训练的效果。 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 参考文献\n",
    "\n",
    "[1] Bennett, Charles H., et al. \"Teleporting an unknown quantum state via dual classical and Einstein-Podolsky-Rosen channels.\" [Physical Review Letters 70.13 (1993): 1895.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.70.1895)\n",
    "\n",
    "[2] Boschi, Danilo, et al. \"Experimental realization of teleporting an unknown pure quantum state via dual classical and Einstein-Podolsky-Rosen channels.\" [Physical Review Letters 80.6 (1998): 1121.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.80.1121)\n",
    "\n",
    "[3] Bouwmeester, Dik, et al. \"Experimental quantum teleportation.\" [Nature 390.6660 (1997): 575-579.](https://www.nature.com/articles/37539)\n",
    "\n",
    "[4] Vidal, Guifre, and Christopher M. Dawson. \"Universal quantum circuit for two-qubit transformations with three controlled-NOT gates.\" [Physical Review A 69.1 (2004): 010301.](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.69.010301)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
